#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# Unit tests for ckpt_protocols.py
import ckpt_protocols
import filesystem
import generic_mock
import unittest
import cPickle
import lambda_algorithm
import services
import types

class TestProtocol(unittest.TestCase):
    
    def setUp(self):

        self.clock_ = generic_mock.GenericMock()
        self.clock_.StubForTesting('get_pid', lambda: 1)

        self.mock_fs_ = filesystem.MockFileSystem()
        self.protocol_ = ckpt_protocols.protocol(
            self.clock_, self.mock_fs_)

    def testConstructor(self):
        self.assertEqual(self.protocol_.get_pid(), 1)
        self.assertEqual(self.protocol_.callback_(), None)
        self.assertEqual(self.protocol_.ck_counter_, -1)
        self.assertEqual(self.protocol_.fs_, self.mock_fs_)

    def testPrintOutput(self):
        self.protocol_.print_output('dummy text')
        self.assertEqual(
            self.clock_, 
            [('print_output', ('dummy text',), {})])

    def testCloseChannel(self):
        self.protocol_.close_channel(2)
        self.assertEqual(
            self.clock_, 
            [('close_channel', (2,), {})])

    def testSendMessage(self):
        self.protocol_.send_message(2, 'msg')
        self.assertEqual(
            self.clock_, 
            [('send_message', (2, 'msg'), {})])

    def testRecvMessage(self):
        self.clock_.StubForTesting('recv_message',
                                   lambda: 'message')
        self.assertEqual(self.protocol_.recv_message(), 
                         'message')
    
    def testNetworkSize(self):
        self.clock_.StubForTesting('network_size',
                                   lambda: 5)
        self.assertEqual(self.protocol_.network_size(), 5) 
        
    def testGetPid(self):
        self.clock_.StubForTesting('get_pid',
                                   lambda: 1)
        self.assertEqual(self.protocol_.get_pid(), 1)

    def testGetHosts(self):
        self.clock_.StubForTesting('get_hosts',
                                   lambda: {1: ('127.0.0.1', 1)})
        self.assertEqual(self.protocol_.get_hosts(), {1: ('127.0.0.1', 1)})
        
    def testChannelState(self):
        self.clock_.StubForTesting('channel_state',
                                   lambda x: True)
        self.assertEqual(self.protocol_.channel_state(1), True)

    def testClose(self):
        self.protocol_.close()
        self.assertEqual(self.clock_, [('close', (), {})])
        
    def testCloseChannel(self):
        self.protocol_.close_channel(1)
        self.assertEqual(self.clock_, [('close_channel', (1,), {})])

    def testGetCheckpointCounter(self):
        self.protocol_.ck_counter_ = 10
        self.assertEqual(self.protocol_.get_checkpoint_counter(), 10)

    def testGetClock(self):
        self.clock_.StubForTesting('get_clock',
                                   lambda: 100)
        self.assertEqual(self.protocol_.get_clock(), 100)

    def testSetCallback(self):
        func = lambda: 0
        self.protocol_.set_callback(func)
        self.assertEqual(self.protocol_.callback_, func)
        self.assertEqual(self.protocol_.callback_(), 0)

        obj = object()
        mthd = types.MethodType(lambda _: 1, obj, object)
        self.protocol_.set_callback(mthd)
        self.assertEqual(self.protocol_.callback_, mthd)
        self.assertEqual(self.protocol_.callback_(), 1)

        self.assertRaises(TypeError, self.protocol_.set_callback, 'ERROR') 
        

    def testTakeCheckpoint_(self):
        callback = (1, 2, 3)
        pid = 1
        clock = 100
        ckcounter = 10
        self.protocol_.set_callback(lambda: callback)
        self.protocol_.ck_counter_ = ckcounter
        self.clock_.StubForTesting('get_clock', lambda: clock)
        self.clock_.StubForTesting('get_pid', lambda: pid)
        self.assertEqual(self.protocol_.take_checkpoint_(
                self.protocol_.FORCED_),
                         ((pid,
                           ckcounter + 1,
                           self.protocol_.FORCED_,
                           clock),
                          callback,
                          None))
        self.assertEqual(self.clock_, 
                         [('increase_clock', (), {}), 
                          ('get_pid', (), {}), 
                          ('get_clock', (), {})])

    def testTakeBasicCheckpoint(self):
        callback = (1, 2, 3)
        pid = 1
        clock = 100
        ckcounter = 10
        self.protocol_.set_callback(lambda: callback)
        self.protocol_.ck_counter_ = ckcounter
        self.clock_.StubForTesting('get_clock', lambda: clock)
        self.clock_.StubForTesting('get_pid', lambda: pid)
        self.assertEqual(self.protocol_.take_basic_checkpoint(),
                         ((pid,
                           ckcounter + 1,
                           self.protocol_.BASIC_,
                           clock),
                          callback,
                          None))

    def testTakeForcedCheckpoint(self):
        callback = (1, 2, 3)
        pid = 1
        clock = 100
        ckcounter = 10
        self.protocol_.set_callback(lambda: callback)
        self.protocol_.ck_counter_ = ckcounter
        self.clock_.StubForTesting('get_clock', lambda: clock)
        self.clock_.StubForTesting('get_pid', lambda: pid)
        self.assertEqual(self.protocol_.take_forced_checkpoint_(),
                         ((pid,
                           ckcounter + 1,
                           self.protocol_.FORCED_,
                           clock),
                          callback,
                          None))

    def testStoreCheckpoint(self):
        checkpoint = ((1, 11, self.protocol_.FORCED_, 100),
                      (1, 2, 3), None)
        self.protocol_.store_checkpoint_(checkpoint)
        self.assertEqual(self.mock_fs_.files_[
                '0001-00000011.ckpt']['content'],
                         cPickle.dumps(checkpoint))

    def testRemoveCheckpoint(self):
        checkpoint = ((1,11),)
        self.protocol_.store_checkpoint_(checkpoint)
        self.assertTrue('0001-00000011.ckpt' in 
                        self.mock_fs_.files_)
        self.protocol_.remove_checkpoint_(1,11)
        self.assertTrue('0001-00000011.ckpt' not in 
                        self.mock_fs_.files_)

class TestBCS(unittest.TestCase):
    
    def setUp(self):

        self.clock_ = generic_mock.GenericMock()
        self.clock_.StubForTesting('get_pid', lambda: 1)

        self.mock_fs_ = filesystem.MockFileSystem()
        self.protocol_ = ckpt_protocols.BCS(
            self.clock_, self.mock_fs_)

    def testConstructor(self):
        self.assertEqual(self.protocol_.lc_, 0)

    def testSendMessage(self):
        self.protocol_.send_message(2, 'msg')
        self.assertEqual(
            self.clock_, 
            [('send_message', (2, ('msg', 1)), {})])

    def testRecvMessage(self):
        self.protocol_.lc_ = 10
        self.clock_.StubForTesting('recv_message',
                                   lambda: (2, ('message', 5)))
        self.assertEqual(self.protocol_.recv_message(), 
                         (2, 'message'))
        self.assertEqual(self.protocol_.lc_, 11)
        self.assertEqual(self.mock_fs_.files_, {})

    def testRecvMessageForced(self):
        self.protocol_.lc_ = 10
        self.protocol_.ck_counter_ = 9
        self.clock_.StubForTesting('get_clock', lambda: 100)
        self.protocol_.set_callback(lambda: (20, 30))
        self.clock_.StubForTesting('recv_message',
                                   lambda: (2, ('message', 15)))
        self.assertEqual(self.protocol_.recv_message(), 
                         (2, 'message'))
        self.assertEqual(self.protocol_.lc_, 15)
        self.assertTrue('0001-00000010.ckpt' in self.mock_fs_.files_)
        obj = cPickle.loads(
            self.mock_fs_.files_['0001-00000010.ckpt']['content'])
        self.assertEqual(obj, ((1, 10, 1, 100), (20, 30), 15))

    def testControlInformation(self):
        self.protocol_.lc_ = 19
        self.assertEqual(self.protocol_.control_information_(), 19)

    def testTakeBasicCheckpoint(self):
        self.protocol_.lc_ = 47
        self.protocol_.ck_counter_ = 65
        self.clock_.StubForTesting('get_clock', lambda: 165)
        self.protocol_.set_callback(lambda: (50, 3))
        self.protocol_.take_basic_checkpoint()
        self.assertEqual(self.protocol_.lc_, 48)
        self.assertTrue('0001-00000066.ckpt' in self.mock_fs_.files_)
        obj = cPickle.loads(
            self.mock_fs_.files_['0001-00000066.ckpt']['content'])
        self.assertEqual(obj, ((1, 66, 0, 165), (50, 3), 48))


class TestFDAS(unittest.TestCase):
    
    def setUp(self):

        self.clock_ = generic_mock.GenericMock()
        self.clock_.StubForTesting('get_pid', lambda: 1)
        self.clock_.StubForTesting('get_hosts', lambda: [1,2,3,4])

        self.mock_fs_ = filesystem.MockFileSystem()
        self.protocol_ = ckpt_protocols.FDAS(
            self.clock_, self.mock_fs_)
        self.clock_.clear()

    def testConstructor(self):
        self.assertEqual(self.protocol_.sent_, False)
        self.assertEqual(self.protocol_.dv_, 
                         {1: 1, 2: 0, 3: 0, 4: 0})

    def testSendMessage(self):
        self.assertEqual(self.protocol_.sent_, False)
        self.protocol_.send_message(2, 'msg')
        self.assertEqual(
            self.clock_,
            [('send_message', (2, ('msg', {1: 1, 2: 0, 3: 0, 4: 0})), {})])
        self.assertEqual(self.protocol_.sent_, True)

    def testRecvMessage(self):
        self.protocol_.dv_ = {1: 3, 2: 5, 3: 2, 4: 0}
        self.clock_.StubForTesting('recv_message',
                                   lambda: (2, ('message', 
                                                {1: 2, 2: 3, 3: 3, 4: 2})))
        self.assertEqual(self.protocol_.recv_message(), 
                         (2, 'message'))
        self.assertEqual(self.protocol_.dv_, {1: 3, 2: 5, 3: 3, 4: 2})
        self.assertEqual(self.mock_fs_.files_, {})

    def testRecvMessageForcedNoMessageSent(self):
        self.protocol_.dv_ = {1: 3, 2: 5, 3: 2, 4: 0}
        self.clock_.StubForTesting('recv_message',
                                   lambda: (2, ('message', 
                                                {1: 2, 2: 7, 3: 3, 4: 2})))
        self.assertEqual(self.protocol_.recv_message(), 
                         (2, 'message'))
        self.assertEqual(self.protocol_.dv_, {1: 3, 2: 7, 3: 3, 4: 2})
        self.assertEqual(self.mock_fs_.files_, {})
        self.assertEqual(self.protocol_.sent_, False)

    def testRecvMessageForced(self):
        self.protocol_.dv_ = {1: 3, 2: 5, 3: 5, 4: 0}
        self.protocol_.ck_counter_ = 3
        self.protocol_.sent_ = True
        self.clock_.StubForTesting('get_clock', lambda: 100)
        self.protocol_.set_callback(lambda: (20, 30))
        self.clock_.StubForTesting('recv_message',
                                   lambda: (2, ('message', 
                                                {1: 2, 2: 7, 3: 5, 4: 2})))
        self.assertEqual(self.protocol_.recv_message(), 
                         (2, 'message'))
        self.assertEqual(self.protocol_.dv_, {1: 4, 2: 7, 3: 5, 4: 2})
        self.assertEqual(self.protocol_.sent_, False)
        self.assertTrue('0001-00000004.ckpt' in self.mock_fs_.files_)
        obj = cPickle.loads(
            self.mock_fs_.files_['0001-00000004.ckpt']['content'])
        self.assertEqual(obj, 
                         ((1, 4, 1, 100), (20, 30), 
                          {1: 3, 2: 5, 3: 5, 4: 0}))
        

    def testControlInformation(self):
        self.protocol_.dv_ = {1: 2, 2: 7, 3: 3, 4: 2}
        self.assertEqual(self.protocol_.control_information_(), 
                         {1: 2, 2: 7, 3: 3, 4: 2})

    def testTakeBasicCheckpoint(self):
        self.protocol_.dv_ = {1: 65, 2: 5, 3: 5, 4: 0}
        self.protocol_.ck_counter_ = 65
        self.protocol_.sent = True
        self.clock_.StubForTesting('get_clock', lambda: 165)
        self.protocol_.set_callback(lambda: (50, 3))
        self.protocol_.take_basic_checkpoint()
        self.assertEqual(self.protocol_.dv_, 
                         {1: 66, 2: 5, 3: 5, 4: 0})
        self.assertEqual(self.protocol_.sent_, False)
        self.assertTrue('0001-00000066.ckpt' in self.mock_fs_.files_)
        obj = cPickle.loads(
            self.mock_fs_.files_['0001-00000066.ckpt']['content'])
        self.assertEqual(obj, ((1, 66, 0, 165), (50, 3), 
                               {1: 65, 2: 5, 3: 5, 4: 0}))


class TestFDASRDTLGC(unittest.TestCase):
    
    def setUp(self):

        self.clock_ = generic_mock.GenericMock()
        self.clock_.StubForTesting('get_pid', lambda: 1)
        self.clock_.StubForTesting('get_hosts', lambda: [1,2,3,4])

        self.mock_fs_ = filesystem.MockFileSystem()
        # It is too simple to be stubbed :-D
        self.protocol_ = ckpt_protocols.RDT_LGC_FDAS(
            self.clock_, self.mock_fs_)
        self.protocol_.take_basic_checkpoint()
        self.clock_.clear()

    def testConstructor(self):
        self.assertEqual(self.protocol_.sent_, False)
        self.assertEqual(self.protocol_.dv_, 
                         {1: 2, 2: 0, 3: 0, 4: 0})
        self.assertEqual(self.protocol_.cm_,
                         {1: [0, 1], 2: None, 3: None, 4: None})

    def testSendMessage(self):
        self.assertEqual(self.protocol_.sent_, False)
        self.protocol_.send_message(2, 'msg')
        self.assertEqual(
            self.clock_,
            [('send_message', (2, ('msg', {1: 2, 2: 0, 3: 0, 4: 0})), {})])
        self.assertEqual(self.protocol_.sent_, True)

    def testRecvMessage(self):
        self.protocol_.dv_ = {1: 3, 2: 5, 3: 2, 4: 0}
        self.clock_.StubForTesting('recv_message',
                                   lambda: (2, ('message', 
                                                {1: 2, 2: 3, 3: 3, 4: 2})))
        self.assertEqual(self.protocol_.recv_message(), 
                         (2, 'message'))
        self.assertEqual(self.protocol_.dv_, {1: 3, 2: 5, 3: 3, 4: 2})
        self.assertEqual(len(self.mock_fs_.files_), 1)
        self.assertEqual(self.protocol_.cm_,
                         {1: [0, 3], 2: None, 
                          3: [0, 3], 4: [0, 3]})


    def testRecvMessageForcedNoMessageSent(self):
        self.protocol_.dv_ = {1: 3, 2: 5, 3: 2, 4: 0}
        self.clock_.StubForTesting('recv_message',
                                   lambda: (2, ('message', 
                                                {1: 2, 2: 7, 3: 3, 4: 2})))
        self.assertEqual(self.protocol_.recv_message(), 
                         (2, 'message'))
        self.assertEqual(self.protocol_.dv_, {1: 3, 2: 7, 3: 3, 4: 2})
        self.assertEqual(len(self.mock_fs_.files_), 1)
        self.assertEqual(self.protocol_.sent_, False)
        self.assertEqual(self.protocol_.cm_,
                         {1: [0, 4], 2: [0, 4], 
                          3: [0, 4], 4: [0, 4]})


    def testRecvMessageForced(self):
        self.protocol_.dv_ = {1: 3, 2: 5, 3: 5, 4: 0}
        self.protocol_.ck_counter_ = 3
        self.protocol_.sent_ = True
        self.protocol_.link_(2, 1)
        self.protocol_.link_(3, 1)
        self.clock_.StubForTesting('get_clock', lambda: 100)
        self.protocol_.set_callback(lambda: (20, 30))
        self.clock_.StubForTesting('recv_message',
                                   lambda: (2, ('message', 
                                                {1: 2, 2: 7, 3: 5, 4: 2})))
        self.assertEqual(self.protocol_.recv_message(), 
                         (2, 'message'))
        self.assertEqual(self.protocol_.dv_, {1: 4, 2: 7, 3: 5, 4: 2})
        self.assertEqual(self.protocol_.sent_, False)
        self.assertEqual(len(self.mock_fs_.files_), 2)
        self.assertTrue('0001-00000004.ckpt' in self.mock_fs_.files_)
        self.assertTrue('0001-00000000.ckpt' in self.mock_fs_.files_)
        obj = cPickle.loads(
            self.mock_fs_.files_['0001-00000004.ckpt']['content'])
        self.assertEqual(obj, 
                         ((1, 4, 1, 100), (20, 30), 
                          {1: 3, 2: 5, 3: 5, 4: 0}))
        self.assertEqual(self.protocol_.cm_,
                         {1: [4, 3], 2: [4, 3], 
                          3: [0, 1], 4: [4, 3]})


    def testRecvMessageForcedRelease(self):
        self.protocol_.dv_ = {1: 3, 2: 5, 3: 5, 4: 0}
        self.protocol_.ck_counter_ = 3
        self.protocol_.sent_ = True
        self.protocol_.link_(2, 1)
        self.protocol_.link_(4, 1)
        self.clock_.StubForTesting('get_clock', lambda: 100)
        self.protocol_.set_callback(lambda: (20, 30))
        self.clock_.StubForTesting('recv_message',
                                   lambda: (2, ('message', 
                                                {1: 2, 2: 7, 3: 5, 4: 2})))
        self.assertEqual(self.protocol_.recv_message(), 
                         (2, 'message'))
        self.assertEqual(self.protocol_.dv_, {1: 4, 2: 7, 3: 5, 4: 2})
        self.assertEqual(self.protocol_.sent_, False)
        # Only P(1) state points to the basic state originally
        self.assertEqual(len(self.mock_fs_.files_), 1)
        self.assertTrue('0001-00000004.ckpt' in self.mock_fs_.files_)
        self.assertTrue('0001-00000000.ckpt' not in self.mock_fs_.files_)
        obj = cPickle.loads(
            self.mock_fs_.files_['0001-00000004.ckpt']['content'])
        self.assertEqual(obj, 
                         ((1, 4, 1, 100), (20, 30), 
                          {1: 3, 2: 5, 3: 5, 4: 0}))
        self.assertEqual(self.protocol_.cm_,
                         {1: [4, 3], 2: [4, 3], 
                          3: None, 4: [4, 3]})

        

    def testControlInformation(self):
        self.protocol_.dv_ = {1: 2, 2: 7, 3: 3, 4: 2}
        self.assertEqual(self.protocol_.control_information_(), 
                         {1: 2, 2: 7, 3: 3, 4: 2})

    def testTakeBasicCheckpoint(self):
        self.protocol_.dv_ = {1: 65, 2: 5, 3: 5, 4: 0}
        self.protocol_.ck_counter_ = 65
        self.protocol_.sent = True
        self.clock_.StubForTesting('get_clock', lambda: 165)
        self.protocol_.set_callback(lambda: (50, 3))
        self.protocol_.take_basic_checkpoint()
        self.assertEqual(len(self.mock_fs_.files_), 1)
        self.assertEqual(self.protocol_.dv_, 
                         {1: 66, 2: 5, 3: 5, 4: 0})
        self.assertEqual(self.protocol_.sent_, False)
        self.assertTrue('0001-00000066.ckpt' in self.mock_fs_.files_)
        obj = cPickle.loads(
            self.mock_fs_.files_['0001-00000066.ckpt']['content'])
        self.assertEqual(obj, ((1, 66, 0, 165), (50, 3), 
                               {1: 65, 2: 5, 3: 5, 4: 0}))
        self.assertEqual(self.protocol_.cm_,
                         {1: [66, 1], 2: None, 
                          3: None, 4: None})

    def testTakeBasicCheckpointLinked(self):
        self.protocol_.dv_ = {1: 65, 2: 5, 3: 5, 4: 0}
        self.protocol_.ck_counter_ = 65
        self.protocol_.sent = True
        self.clock_.StubForTesting('get_clock', lambda: 165)
        self.protocol_.set_callback(lambda: (50, 3))
        self.protocol_.link_(3, 1)
        self.protocol_.take_basic_checkpoint()
        self.assertEqual(len(self.mock_fs_.files_), 2)
        self.assertEqual(self.protocol_.dv_, 
                         {1: 66, 2: 5, 3: 5, 4: 0})
        self.assertEqual(self.protocol_.sent_, False)
        self.assertTrue('0001-00000066.ckpt' in self.mock_fs_.files_)
        self.assertTrue('0001-00000000.ckpt' in self.mock_fs_.files_)
        obj = cPickle.loads(
            self.mock_fs_.files_['0001-00000066.ckpt']['content'])
        self.assertEqual(obj, ((1, 66, 0, 165), (50, 3), 
                               {1: 65, 2: 5, 3: 5, 4: 0}))
        self.assertEqual(self.protocol_.cm_,
                         {1: [66, 1], 2: None, 
                          3: [0, 1], 4: None})


    def testNewCkpt(self):
        self.assertEqual(self.protocol_.cm_,
                         {1: [0, 1], 2: None, 3: None, 4: None})
        self.protocol_.new_ckpt_(1, 20)
        self.assertEqual(self.protocol_.cm_,
                         {1: [20, 1], 2: None, 3: None, 4: None})
        self.protocol_.new_ckpt_(2, 15)
        self.assertEqual(self.protocol_.cm_,
                         {1: [20, 1], 2: [15, 1], 3: None, 4: None})
   
    def testLink(self):
        self.protocol_.link_(2, 1)
        self.assertEqual(self.protocol_.cm_,
                         {1: [0, 2], 2: [0, 2], 3: None, 4: None})

    def testFree(self):
        self.protocol_.link_(2, 1)
        self.protocol_.free_(1)
        self.assertEqual(self.protocol_.cm_,
                         {1: None, 2: [0, 1], 3: None, 4: None})
        f = self.mock_fs_.OpenFile('0001-00000000.ckpt', 'w')
        f.close()
        self.assertTrue('0001-00000000.ckpt' in self.mock_fs_.files_)
        self.protocol_.free_(2)
        self.assertTrue('0001-00000000.ckpt' not in self.mock_fs_.files_)


class TestFDASRDTLESS(unittest.TestCase):
    
    def setUp(self):

        self.clock_ = generic_mock.GenericMock()
        self.clock_.StubForTesting('get_pid', lambda: 1)
        self.clock_.StubForTesting('get_hosts', lambda: [1,2,3,4])

        self.mock_fs_ = filesystem.MockFileSystem()
        self.alg_ = generic_mock.GenericMock()
        self.alg_.StubForTesting('packages', lambda: 2)
        self.alg_.StubForTesting('Disperse', lambda x: ['A', 'B'])
        self.ring_ = generic_mock.GenericMock()
        
        self.protocol_ = ckpt_protocols.RDT_DISKLESS_FDAS(
            self.clock_, self.mock_fs_, self.alg_, self.ring_)
        self.protocol_.take_basic_checkpoint()
        self.clock_.clear()

    def testConstructor(self):
        self.assertEqual(self.protocol_.sent_, False)
        self.assertEqual(self.protocol_.dv_, 
                         {1: 2, 2: 0, 3: 0, 4: 0})
        self.assertEqual(self.protocol_.cm_,
                         {1: [0, 1, None], 2: None, 3: None, 4: None})
        self.assertEqual(self.ring_, [('set_comm', ([],), {}), 
                                      ('send_pos', (0, [1, 0, 'A']), {})])

    def testSendMessage(self):
        self.assertEqual(self.protocol_.sent_, False)
        self.protocol_.send_message(2, 'msg')
        self.assertEqual(
            self.clock_,
            [('send_message', 
              (2, [0, 
                   ('msg', {1: 2, 2: 0, 3: 0, 4: 0}, 
                    {1: -1, 2: -1, 3: -1, 4: -1})]), {})])
        self.assertEqual(self.protocol_.sent_, True)

    def testRecvMessage(self):
        self.ring_.clear()
        self.protocol_.dv_ = {1: 3, 2: 5, 3: 2, 4: 0}
        self.clock_.StubForTesting(
            'recv_message',
            lambda: (2, (self.protocol_.MESSAGE_,
                    ('message', 
                     {1: 2, 2: 3, 3: 3, 4: 2},
                     {1: -1, 2: 100, 3: -1, 4: -1}))))
        self.assertEqual(self.protocol_.recv_message(), 
                         (2, 'message'))
        self.assertEqual(self.protocol_.dv_, {1: 3, 2: 5, 3: 3, 4: 2})
        self.assertEqual(len(self.mock_fs_.files_), 2)
        self.assertEqual(self.protocol_.cm_,
                         {1: [0, 3, None], 2: None, 
                          3: [0, 3, None], 4: [0, 3, None]})
        self.assertEqual(self.ring_, [])
        self.assertEqual(self.protocol_.des_, [])
        self.assertEqual(self.protocol_.dvk_, 
                         {1: -1, 2: 100, 3: -1, 4: -1})
        
    def testRecvMessageForcedNoMessageSent(self):
        self.ring_.clear()
        self.protocol_.dv_ = {1: 3, 2: 5, 3: 2, 4: 0}
        self.clock_.StubForTesting(
            'recv_message',
            lambda: (2, (self.protocol_.MESSAGE_,
                    ('message', 
                     {1: 2, 2: 7, 3: 3, 4: 2},
                     {1: 0, 2: 100, 3: 0, 4: 0}))))
        self.assertEqual(self.protocol_.recv_message(), 
                         (2, 'message'))
        self.assertEqual(self.protocol_.dv_, {1: 3, 2: 7, 3: 3, 4: 2})
        self.assertEqual(len(self.mock_fs_.files_), 2)
        self.assertEqual(self.protocol_.sent_, False)
        self.assertEqual(self.protocol_.cm_,
                         {1: [0, 4, None], 2: [0, 4, None], 
                          3: [0, 4, None], 4: [0, 4, None]})
        self.assertEqual(self.protocol_.des_, [])
        self.assertEqual(self.ring_, [])
        self.assertEqual(self.protocol_.dvk_, 
                         {1: 0, 2: 100, 3: 0, 4: 0})

    def testRecvMessageForced(self):
        self.ring_.clear()
        self.protocol_.dv_ = {1: 3, 2: 5, 3: 5, 4: 0}
        self.protocol_.ck_counter_ = 3
        self.protocol_.sent_ = True
        self.protocol_.link_(2, 1)
        self.protocol_.link_(3, 1)
        self.clock_.StubForTesting('get_clock', lambda: 100)
        self.protocol_.set_callback(lambda: (20, 30))
        self.clock_.StubForTesting(
            'recv_message',
            lambda: (2, (self.protocol_.MESSAGE_,
                    ('message', 
                     {1: 2, 2: 7, 3: 5, 4: 2},
                     {1: 0, 2: 100, 3: 0, 4: 0}))))
        self.assertEqual(self.protocol_.recv_message(), 
                         (2, 'message'))
        self.assertEqual(self.protocol_.dv_, {1: 4, 2: 7, 3: 5, 4: 2})
        self.assertEqual(self.protocol_.sent_, False)
        self.assertEqual(len(self.mock_fs_.files_), 4)
        self.assertTrue('0001-00000004.ckpt' in self.mock_fs_.files_)
        self.assertTrue('0001-00000004.package' in self.mock_fs_.files_)
        obj = cPickle.loads(
            self.mock_fs_.files_['0001-00000004.ckpt']['content'])
        self.assertEqual(obj, 
                         ((1, 4, 1, 100), (20, 30), 
                          {1: 3, 2: 5, 3: 5, 4: 0}))
        self.assertEqual(self.protocol_.cm_,
                         {1: [4, 3, None], 2: [4, 3, None], 
                          3: [0, 1, None], 4: [4, 3, None]})
        self.assertEqual(self.protocol_.des_, [])
        self.assertEqual(self.ring_,
                         [('send_pos', (0, [1, 4, 'A']), {})])
        self.assertEqual(self.protocol_.dvk_, 
                         {1: 0, 2: 100, 3: 0, 4: 0})

    def testRecvMessageForcedRelease(self):
        self.ring_.clear()
        self.protocol_.dv_ = {1: 3, 2: 5, 3: 5, 4: 0}
        self.protocol_.ck_counter_ = 3
        self.protocol_.sent_ = True
        self.protocol_.link_(2, 1)
        self.protocol_.link_(4, 1)
        self.clock_.StubForTesting('get_clock', lambda: 100)
        self.protocol_.set_callback(lambda: (20, 30))
        self.clock_.StubForTesting(
            'recv_message',
            lambda: (2, (self.protocol_.MESSAGE_,
                    ('message', 
                     {1: 2, 2: 7, 3: 5, 4: 2},
                     {1: 0, 2: 0, 3: 100, 4: 40}))))
        self.assertEqual(self.protocol_.recv_message(), 
                         (2, 'message'))
        self.assertEqual(self.protocol_.dv_, {1: 4, 2: 7, 3: 5, 4: 2})
        self.assertEqual(self.protocol_.sent_, False)
        self.assertEqual(len(self.mock_fs_.files_), 4)
        self.assertTrue('0001-00000004.ckpt' in self.mock_fs_.files_)
        self.assertTrue('0001-00000004.package' in self.mock_fs_.files_)
        obj = cPickle.loads(
            self.mock_fs_.files_['0001-00000004.ckpt']['content'])
        self.assertEqual(obj, 
                         ((1, 4, 1, 100), (20, 30), 
                          {1: 3, 2: 5, 3: 5, 4: 0}))
        self.assertEqual(self.protocol_.cm_,
                         {1: [4, 3, None], 2: [4, 3, None], 
                          3: None, 4: [4, 3, None]})
        self.assertEqual(self.protocol_.des_,
                         [[0, 0, {1: 4, 2: 7, 3: 5, 4: 2}]])
        self.assertEqual(self.ring_,
                         [('send_pos', (0, [1, 4, 'A']), {})])
        self.assertEqual(self.protocol_.dvk_, 
                         {1: 0, 2: 0, 3: 100, 4: 40})

    def testControlInformation(self):
        self.protocol_.dv_ = {1: 2, 2: 7, 3: 3, 4: 2}
        self.assertEqual(self.protocol_.control_information_(), 
                         {1: 2, 2: 7, 3: 3, 4: 2})

    def testTakeBasicCheckpoint(self):
        self.ring_.clear()
        self.protocol_.dv_ = {1: 65, 2: 5, 3: 5, 4: 0}
        self.protocol_.dvk_ = self.protocol_.dv_
        self.protocol_.ck_counter_ = 65
        self.protocol_.sent = True
        self.clock_.StubForTesting('get_clock', lambda: 165)
        self.protocol_.set_callback(lambda: (50, 3))
        self.protocol_.take_basic_checkpoint()
        self.assertEqual(len(self.mock_fs_.files_), 3)
        self.assertEqual(self.protocol_.dv_, 
                         {1: 66, 2: 5, 3: 5, 4: 0})
        self.assertEqual(self.protocol_.sent_, False)
        self.assertTrue('0001-00000066.ckpt' in self.mock_fs_.files_)
        self.assertTrue('0001-00000066.package' in self.mock_fs_.files_)
        self.assertFalse('0001-00000000.package' in self.mock_fs_.files_)
        obj = cPickle.loads(
            self.mock_fs_.files_['0001-00000066.ckpt']['content'])
        self.assertEqual(obj, ((1, 66, 0, 165), (50, 3), 
                               {1: 65, 2: 5, 3: 5, 4: 0}))
        self.assertEqual(
            self.mock_fs_.files_['0001-00000066.package']['content'],
            'B')
        self.assertEqual(self.protocol_.cm_,
                         {1: [66, 1, None], 2: None, 
                          3: None, 4: None})
        self.assertEqual(
            self.ring_, 
            [('send_pos', (0, [1, 66, 'A']), {}), 
             ('send_pos', (0, [3, 0]), {})])
        self.assertEqual(self.protocol_.des_, [])

    def testTakeBasicCheckpointNotPackageUnecessary(self):
        self.ring_.clear()
        self.protocol_.dv_ = {1: 65, 2: 5, 3: 5, 4: 0}
        self.protocol_.ck_counter_ = 65
        self.protocol_.sent = True
        self.clock_.StubForTesting('get_clock', lambda: 165)
        self.protocol_.set_callback(lambda: (50, 3))
        self.protocol_.take_basic_checkpoint()
        self.assertEqual(len(self.mock_fs_.files_), 4)
        self.assertEqual(self.protocol_.dv_, 
                         {1: 66, 2: 5, 3: 5, 4: 0})
        self.assertEqual(self.protocol_.sent_, False)
        self.assertTrue('0001-00000066.ckpt' in self.mock_fs_.files_)
        self.assertTrue('0001-00000066.package' in self.mock_fs_.files_)
        self.assertTrue('0001-00000000.package' in self.mock_fs_.files_)
        obj = cPickle.loads(
            self.mock_fs_.files_['0001-00000066.ckpt']['content'])
        self.assertEqual(obj, ((1, 66, 0, 165), (50, 3), 
                               {1: 65, 2: 5, 3: 5, 4: 0}))
        self.assertEqual(
            self.mock_fs_.files_['0001-00000066.package']['content'],
            'B')
        self.assertEqual(self.protocol_.cm_,
                         {1: [66, 1, None], 2: None, 
                          3: None, 4: None})
        self.assertEqual(
            self.ring_, 
            [('send_pos', (0, [1, 66, 'A']), {})])
        self.assertEqual(self.protocol_.des_, 
                         [[0, 0, {1: 65, 2: 5, 3: 5, 4: 0}]])

    def testTakeBasicCheckpointLinked(self):
        self.ring_.clear()
        self.protocol_.dv_ = {1: 65, 2: 5, 3: 5, 4: 0}
        self.protocol_.dvk_ = self.protocol_.dv_
        self.protocol_.ck_counter_ = 65
        self.protocol_.sent = True
        self.clock_.StubForTesting('get_clock', lambda: 165)
        self.protocol_.set_callback(lambda: (50, 3))
        self.protocol_.link_(3, 1)
        self.protocol_.take_basic_checkpoint()
        self.assertEqual(len(self.mock_fs_.files_), 4)
        self.assertEqual(self.protocol_.dv_, 
                         {1: 66, 2: 5, 3: 5, 4: 0})
        self.assertEqual(self.protocol_.sent_, False)
        self.assertTrue('0001-00000066.ckpt' in self.mock_fs_.files_)
        self.assertTrue('0001-00000066.package' in self.mock_fs_.files_)
        self.assertTrue('0001-00000000.package' in self.mock_fs_.files_)
        obj = cPickle.loads(
            self.mock_fs_.files_['0001-00000066.ckpt']['content'])
        self.assertEqual(obj, ((1, 66, 0, 165), (50, 3), 
                               {1: 65, 2: 5, 3: 5, 4: 0}))
        self.assertEqual(self.protocol_.cm_,
                         {1: [66, 1, None], 2: None, 
                          3: [0, 1, None], 4: None})
        self.assertEqual(self.ring_, [('send_pos', (0, [1, 66, 'A']), {})])
        self.assertEqual(self.protocol_.des_, [])

    def testNewCkpt(self):
        self.assertEqual(self.protocol_.cm_,
                         {1: [0, 1, None], 2: None, 3: None, 4: None})
        self.protocol_.new_ckpt_(1, 20)
        self.assertEqual(self.protocol_.cm_,
                         {1: [20, 1, None], 2: None, 3: None, 4: None})
        self.protocol_.new_ckpt_(2, 15)
        self.assertEqual(self.protocol_.cm_,
                         {1: [20, 1, None], 2: [15, 1, None], 
                          3: None, 4: None})
   
    def testLink(self):
        self.protocol_.link_(2, 1)
        self.assertEqual(self.protocol_.cm_,
                         {1: [0, 2, None], 2: [0, 2, None], 
                          3: None, 4: None})

    def testFree(self):
        self.protocol_.link_(2, 1)
        self.protocol_.free_(1)
        self.assertEqual(self.protocol_.cm_,
                         {1: None, 2: [0, 1, None], 3: None, 4: None})
        self.protocol_.free_(2)
        self.assertEqual(self.protocol_.cm_,
                 {1: None, 2: None, 3: None, 4: None})
        self.assertEqual(self.protocol_.des_, 
                         [[0, 0, {1: 2, 2: 0, 3: 0, 4: 0}]])

    def testCollect(self):
        self.ring_.clear()
        self.protocol_.dvk_ = {1: 2, 2: 0, 3: 0, 4: 0}
        self.protocol_.des_.append([0, 0, {1: 2, 2: 0, 3: 0, 4: 0}])
        self.protocol_.des_.append([5, 0, {1: 10, 2: 0, 3: 0, 4: 0}])
        self.assertTrue('0001-00000000.package' in self.mock_fs_.files_)
        self.protocol_.collect_()
        self.assertEqual(self.protocol_.des_, 
                         [[5, 0, {1: 10, 2: 0, 3: 0, 4: 0}]])
        self.assertFalse('0001-00000000.package' in self.mock_fs_.files_)
        self.assertEqual(self.ring_, [('send_pos', (0, [3, 0]), {})])

    def testPackageUnecessary(self):
        self.protocol_.dvk_ = {1: 2, 2: 0, 3: 0, 4: 0}
        ckpt = [0, 0, {1: 2, 2: 0, 3: 1, 4: 0}]
        self.assertFalse(self.protocol_.package_unecessary_(ckpt))
        ckpt = [0, 0, {1: 3, 2: 0, 3: 0, 4: 0}]
        self.assertFalse(self.protocol_.package_unecessary_(ckpt))
        ckpt = [0, 0, {1: 2, 2: 0, 3: 0, 4: 0}]
        self.assertTrue(self.protocol_.package_unecessary_(ckpt))
        ckpt = [0, 0, {1: 0, 2: 0, 3: 0, 4: 0}]
        self.assertTrue(self.protocol_.package_unecessary_(ckpt))

    def testStorePackage(self):
        self.protocol_.store_package_(5, 3, 'DATA')
        self.assertTrue('0005-00000003.package' in self.mock_fs_.files_)
        self.assertEqual(
            self.mock_fs_.files_['0005-00000003.package']['content'], 
            'DATA')

    def testRemovePackage(self):
        self.protocol_.store_package_(5, 3, 'DATA')
        self.assertTrue('0005-00000003.package' in self.mock_fs_.files_)
        self.protocol_.remove_package_(5, 3)
        self.assertFalse('0005-00000003.package' in self.mock_fs_.files_)

    def testRecvStoreMessage(self):
        msg = [(3, (self.protocol_.STORE_, 5, 'AAA')), 
               (2, (self.protocol_.MESSAGE_,
                    ('message', 
                     {1: 0, 2: 0, 3: 0, 4: 0},
                     {1: 0, 2: 0, 3: 0, 4: 0})))]
        self.clock_.StubForTesting(
            'recv_message', lambda: msg.pop(0))
        self.protocol_.recv_message()
        self.assertTrue('0003-00000005.package' in self.mock_fs_.files_)
        self.assertEqual(
            self.mock_fs_.files_['0003-00000005.package']['content'], 
            'AAA')
        self.assertEqual(self.clock_, 
                         [('recv_message', (), {}), 
                          ('send_message', (3, [2, 5]), {}), 
                          ('recv_message', (), {}), 
                          ('get_hosts', (), {})])

    def testRecvConfirmMessage(self):
        self.alg_.StubForTesting('packages', lambda: 3)
        msg = [(3, (self.protocol_.CONFIRM_, 5)), 
               (2, (self.protocol_.MESSAGE_,
                    ('message', 
                     {1: 0, 2: 0, 3: 0, 4: 0},
                     {1: -1, 2: -1, 3: -1, 4: -1})))]
        self.clock_.StubForTesting(
            'recv_message', lambda: msg.pop(0))
        self.protocol_.recv_message()
        self.assertEqual(self.protocol_.max_, {1: 0, 3: 5})
        self.assertEqual(self.protocol_.lc_, -1)
        self.assertTrue('0001-00000000.ckpt' in self.mock_fs_.files_)
        msg = [(2, (self.protocol_.CONFIRM_, 3)), 
               (2, (self.protocol_.MESSAGE_,
                    ('message', 
                     {1: 0, 2: 0, 3: 0, 4: 0},
                     {1: -1, 2: -1, 3: -1, 4: -1})))]
        self.protocol_.recv_message()
        self.assertEqual(self.protocol_.max_, {1: 0, 2: 3, 3: 5})
        self.assertEqual(self.protocol_.lc_, 0)
        self.assertFalse('0001-00000000.ckpt' in self.mock_fs_.files_)
        
    def testRecvEraseMessage(self):
        msg = [(3, (self.protocol_.ERASE_, 0)), 
               (2, (self.protocol_.MESSAGE_,
                    ('message', 
                     {1: 0, 2: 0, 3: 0, 4: 0},
                     {1: 0, 2: 0, 3: 0, 4: 0})))]
        self.clock_.StubForTesting(
            'recv_message', lambda: msg.pop(0))
        self.protocol_.store_package_(3, 0, 'AAA')
        self.assertTrue('0003-00000000.package' in self.mock_fs_.files_)
        self.protocol_.recv_message()
        self.assertFalse('0003-00000000.package' in self.mock_fs_.files_)

# TODO(menderico): Test propagate class

if __name__ == '__main__':
    unittest.main()
